import sys
sys.path.append("..")
import pyrtl
from  pyrtl import GPUSim
import or1200_definitions

class Or1200Pm(object):
    def __init__(self):
        self.rst = pyrtl.Input(bitwidth=1, name='Or1200Pm_rst')
        self.pic_wakeup = pyrtl.Input(bitwidth=1, name='Or1200Pm_pic_wakeup')
        self.spr_write = pyrtl.Input(bitwidth=1, name='Or1200Pm_spr_write')
        self.spr_addr = pyrtl.Input(bitwidth=32, name='Or1200Pm_spr_addr')
        self.spr_dat_i = pyrtl.Input(bitwidth=32, name='Or1200Pm_spr_dat_i')
        self.spr_dat_o = pyrtl.Output(bitwidth=32, name='Or1200Pm_spr_dat_o')
        #self.spr_dat_o_tmp = pyrtl.WireVector(bitwidth=32, name='spr_dat_o_tmp')

        self.pm_cpustall = pyrtl.Input(bitwidth=1, name='Or1200Pm_pm_cpustall')
        self.pm_clksd = pyrtl.Output(bitwidth=4, name='Or1200Pm_pm_clksd')
        self.pm_dc_gate = pyrtl.Output(bitwidth=1, name='Or1200Pm_pm_dc_gate')
        self.pm_ic_gate = pyrtl.Output(bitwidth=1, name='Or1200Pm_pm_ic_gate')
        self.pm_dmmu_gate = pyrtl.Output(bitwidth=1, name='Or1200Pm_pm_dmmu_gate')
        self.pm_immu_gate = pyrtl.Output(bitwidth=1, name='Or1200Pm_pm_immu_gate')
        self.pm_tt_gate = pyrtl.Output(bitwidth=1, name='Or1200Pm_pm_tt_gate')
        self.pm_cpu_gate = pyrtl.Output(bitwidth=1, name='Or1200Pm_pm_cpu_gate')
        self.pm_wakeup = pyrtl.Output(bitwidth=1, name='Or1200Pm_pm_wakeup')
        self.pm_lvolt = pyrtl.Output(bitwidth=1, name='Or1200Pm_pm_lvolt')

        self.sdf = pyrtl.Register(bitwidth=4, name='Or1200Pm_sdf')
        self.dme = pyrtl.Register(bitwidth=1, name='Or1200Pm_dme')
        self.sme = pyrtl.Register(bitwidth=1, name='Or1200Pm_sme')
        self.dcge = pyrtl.Register(bitwidth=1, name='Or1200Pm_dcge')

        self.pmr_sel = pyrtl.WireVector(bitwidth=1, name='Or1200Pm_pmr_sel')

        self.pmr_sel <<= pyrtl.select(self.spr_addr[or1200_definitions.OR1200_SPR_GROUP_BITS].__eq__(or1200_definitions.OR1200_SPRGRP_PM),\
                                      pyrtl.Const(0b1, bitwidth=1), pyrtl.Const(0b0, bitwidth=1))

        # initialize all the parts
        write2Pmr = Write2Pmr()

        # establish connection relations for Write to PMR
        self.dcge.next <<= write2Pmr.dcge
        self.sdf.next <<= write2Pmr.sdf
        self.sme.next <<= write2Pmr.sme
        self.dme.next <<= write2Pmr.dme
        write2Pmr.rst <<= self.rst
        write2Pmr.pmr_sel <<= self.pmr_sel
        write2Pmr.spr_write <<= self.spr_write
        write2Pmr.pic_wakeup <<= self.pic_wakeup
        write2Pmr.spr_dat_i <<= self.spr_dat_i

        # OR1200_PM_READREGS is True
        #Here or1200_definitions.OR1200_PM_PMR_SDF is a slice from 0 to 4
        #which means select bits 0,1,2,3 from self.spr_dat_o to be assigned
        #bit-by-bit by self.sdf
        self.spr_dat_o <<= pyrtl.concat(self.sdf, self.dme, self.sme, self.dcge, pyrtl.Const(0b0, bitwidth=25))

        # Generate pm_clksd
        self.pm_clksd <<= self.sdf

        # Statically generate all clock gate outputs
        # TODO: add dynamic clock gating feature
        self.pm_cpu_gate <<= (self.dme | self.sme) & ~self.pic_wakeup
        self.pm_cpu_gate_tmp = pyrtl.WireVector(bitwidth=1, name='pm_cpu_gate ')
        self.pm_cpu_gate_tmp <<= (self.dme | self.sme) & ~self.pic_wakeup
        self.pm_dc_gate <<= self.pm_cpu_gate_tmp
        self.pm_ic_gate <<= self.pm_cpu_gate_tmp
        self.pm_dmmu_gate <<= self.pm_cpu_gate_tmp
        self.pm_immu_gate <<= self.pm_cpu_gate_tmp
        self.pm_tt_gate <<= self.sme & ~self.pm_cpu_gate_tmp

        # Assert pm_wakeup when pic_wakeup is asserted
        self.pm_wakeup <<= self.pic_wakeup

        # Assert pm_lvolt when pm_cpu_gate or pm_cpustall are asserted
        self.pm_lvolt <<= self.pm_cpu_gate_tmp | self.pm_cpustall

# Write to PMR and also PMR[DME]/PMR[SME] reset when
# pic_wakeup is asserted
class Write2Pmr(object):
    def __init__(self):
        self.rst = pyrtl.WireVector(bitwidth=1, name='Write2Pmr_rst')
        self.sdf = pyrtl.WireVector(bitwidth=4, name='Write2Pmr_sdf')
        self.dme = pyrtl.WireVector(bitwidth=1, name='Write2Pmr_dme')
        self.sme = pyrtl.WireVector(bitwidth=1, name='Write2Pmr_sme')
        self.dcge = pyrtl.WireVector(bitwidth=1, name='Write2Pmr_dcge')
        self.pmr_sel = pyrtl.WireVector(bitwidth=1, name='Write2Pmr_pmr_sel')
        self.spr_write = pyrtl.WireVector(bitwidth=1, name='Write2Pmr_spr_write')
        self.spr_dat_i = pyrtl.WireVector(bitwidth=32, name='Write2Pmr_spr_dat_i')
        self.pic_wakeup = pyrtl.WireVector(bitwidth=1, name='Write2Pmr_pic_wakeup')

        with pyrtl.conditional_assignment:
            with self.rst:
                self.dcge |= pyrtl.Const(0b0, bitwidth=1)
                self.dme |= pyrtl.Const(0b0, bitwidth=1)
                self.sme |= pyrtl.Const(0b0, bitwidth=1)
                self.sdf |= pyrtl.Const(0b0, bitwidth=4)
            with self.pmr_sel & self.spr_write:
                self.sdf |= self.spr_dat_i[or1200_definitions.OR1200_PM_PMR_SDF]
                self.dme |= self.spr_dat_i[or1200_definitions.OR1200_PM_PMR_DME]
                self.sme |= self.spr_dat_i[or1200_definitions.OR1200_PM_PMR_SME]
                self.dcge |= self.spr_dat_i[or1200_definitions.OR1200_PM_PMR_DCGE]
            with self.pic_wakeup:
                self.dme |= pyrtl.Const(0b0, bitwidth=1)
                self.sme |= pyrtl.Const(0b0, bitwidth=1)


if __name__ == '__main__':
    or1200Pm = Or1200Pm()
    sim = pyrtl.GPUSim_now.GPUSim(65536)
    sim.create_dll('pm.cu')
    print(pyrtl.working_block())
    ff = open("D:\work_doc\PyRTL\examples\Opencore\or1200_pm.c", 'w')
    essent = pyrtl.ESSENT()
    essent._create_code(lambda s: ff.write(s + '\n'))
    ff.close()
    com = pyrtl.CompiledSimulation()

    input_width = []
    from pyrtl.wire import Input

    IN = pyrtl.working_block().wirevector_subset(Input)
    for inn in IN:
        name = inn.name
        width = inn.bitwidth
        input_width.append([name, width])

    import random

    for i in range(10):

        inputs = {}
        inputs_old = {}

        for inp in input_width:
            inputs[inp[0]] = []

            result = random.randint(0, 2 ** inp[1] - 1)
            result_old = result
            inputs_old[inp[0]] = [result, result_old]

        for j in range(10000):
            for inp in input_width:
                if random.randint(0, 20) in [0, 5]:
                    result = random.randint(0, 2 ** inp[1] - 1)
                    inputs[inp[0]].append(result)
                    inputs_old[inp[0]][1] = result
                else:
                    inputs[inp[0]].append(inputs_old[inp[0]][1])

        time1 = essent.step_multiple(inputs)
        time2 = com.step_multiple(inputs)

        print('==========')
        print('essent:', time1)
        print('com:', time2)
        print(time2 / time1)